# %%import
from reflow.utils import _PIPELINES, _SCHEDULERS
from copy import deepcopy
import torch
from tqdm import tqdm, trange
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from torchvision.utils import make_grid
from typing import Callable, List, Optional, Union
from reflow.utils import decode_latents, nothing
from reflow.utils import set_seed
from diffusers.utils import randn_tensor
import math
from diffusers import UNet2DConditionModel
from loguru import logger 
from PIL import Image
from transformers import Blip2Processor, Blip2ForConditionalGeneration


# %%config设置
device = 'cuda:0'
diffusers_pipeline = 'stable_diffusion'
diffusers_scheduler = 'dpm_solver_multi'
diffusers_pipeline_ckpt = 'checkpoints/SD-1-4'
ckpt_path=None
# ckpt_path = "logs/pokemon/distill/init2Reflow_s5000/checkpoints/score_model_s5000.pth"

use_xformers=True

# %%加载 diffusion pipeline
pipeline_cls = _PIPELINES[diffusers_pipeline]
scheduler_cls = _SCHEDULERS[diffusers_scheduler]

weight_dtype=torch.float16
pipeline = pipeline_cls.from_pretrained(
    diffusers_pipeline_ckpt,
    torch_dtype=weight_dtype,
    safety_checker=None,
    requires_safety_checker=False,
)
if ckpt_path:
    pipeline.unet.load_state_dict(torch.load(ckpt_path))
pipeline.scheduler = scheduler_cls.from_config(pipeline.scheduler.config)
pipeline.vae.load_state_dict(torch.load("checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.bin"))
pipeline = pipeline.to(device)

if use_xformers:
    pipeline.enable_xformers_memory_efficient_attention()

#%% load blip model
processor = Blip2Processor.from_pretrained("checkpoints/blip2-flan-t5-xl-coco")
model = Blip2ForConditionalGeneration.from_pretrained("checkpoints/blip2-flan-t5-xl-coco", torch_dtype=weight_dtype)
model = model.to(device)

    
# %% sample config
inference_steps=25
guidance_scale=7.5
# prompts = [
#     'mid shot portrait of a woman in nightclub, in the style of David cronenberg ,scary, weird, high fashion, ID magazine, vogue magazine, homes and garden magazine, surprising, freaky, freak show, realistic, sharp focus, 8k high definition, medium format film photography, photo realistic, insanely detailed, intricate, elegant, art by les edwards and David kostic and stanley lau and artgerm',
#     "cut paper portrait of worf, klingon, intricate, detailed, sharp focus, layered, paper, unreal engine, cgsociety, patrick cabral, kiri ken, greg rutkowski. ",
#     "a portrait a very ordinary person, by Giuseppe Arcimboldo, portrait, fruit, renaissance, anatomically correct, beautiful perfect face, sharp focus, Highly Detailed",
#     "a cinematic portrait elon musk!! as a trap with cat ears, art by lois van baarle and loish! and rossdraws and sam yang and samdoesarts and artgerm and saruei and disney, digital art, highly detailed, intricate, sharp focus, trending on artstation hq, deviantart, unreal engine 5, 4 k uhd image ",
# ]
# prompts = ["Rare Signed 20x24 ""Rue Aubriot, Parisian Street 1975"" Vintage Silver Gelatin Print by Helmut Newton Photography-Global Images Gallery-Global Images",
# "374 best images about Androgyny--The Female Dandy on Pinterest",
# "<p>David Birrel as Badger, Thomas Howes as Ratty and Fra Fee as Mole in <em>The Wind in the Willows</em>.</p><br />© Marc Brenner/Jamie Hendry Productions",
# "A beautiful night sky behind a shining lighthouse. Chania, Crete, Greece",]

# prompts = [
#     "a black and white photo of a street with a light pole",
# "a woman in a suit and tie with her hand on her chin",
# "a group of men in costumes are standing on stage",
# "a lighthouse sitting on a pier with a starry sky",
# ]

# # * pokemon-mix
# prompts = [
#     "Girl with a pearl earring",
#     "Hello Kitty",
#     "Donald Trump",
#     "Totoro",
# ]

# * mostly coco
prompts = [
    'A young man smiles and holds a small teddy bear.',
    'An older man watches a kite fly from across a body of water.',
    "hyperdetailed robotic skeleton head with blue human eyes, symetry, golden ratio, intricate, detailed,",
    "The face and ear of a teddy bear",
    "A double decker bus going down the street.",
    "a meat sandwich that is on french bread.",
    'A kitchen with an oven, stove, cabinets and knives',
    'A man holding a frisbee in a parking lot near water.',
    "A black dragon with red demonic eyes",
]

# %%使用 diffusion model 采样获取数据
images = pipeline(
    prompt=prompts,
    # latents=noise, 
    num_inference_steps=inference_steps,
    guidance_scale=guidance_scale,
).images


# %% auto caption
inputs = processor(images, ["a photo of"]*len(images), return_tensors="pt").to(device, dtype=weight_dtype)
out = model.generate(**inputs)
#%% show outcome
import matplotlib.pyplot as plt

def show_images_in_grid(images):
    """
    根据图像数量自动计算行数和列数，并在网格中展示给定的 PIL 图像列表。
    
    Args:
        images: 包含 PIL Image 对象的列表。
    """
    num_images = len(images)
    cols = math.ceil(math.sqrt(num_images))
    rows = math.ceil(num_images / cols)
    fig = plt.figure(figsize=(cols*5, rows*5))
    for i in range(1, rows * cols + 1):
        if i <= len(images):
            img = images[i - 1]
            fig.add_subplot(rows, cols, i)
            plt.imshow(img)
        else:
            fig.add_subplot(rows, cols, i)
    plt.show()
    
show_images_in_grid(images,)

auto_caption = processor.batch_decode(out, skip_special_tokens=True, clean_up_tokenization_spaces =True)
for caption in auto_caption:
    print(caption)

# %%
